#!/usr/bin/env python

import argparse
import base64
import configparser
import copy
import json
import os
import pickle
import re
import sys
import time
from typing import Dict, Tuple

import requests
from Crypto.Cipher import PKCS1_v1_5
from Crypto.PublicKey import RSA
from prompt_toolkit import prompt
from prompt_toolkit.completion import NestedCompleter
from prompt_toolkit.document import Document
from prompt_toolkit.validation import Validator, ValidationError


def main() -> None:
    """
    程序入口，主要负责处理命令行参数
    """
    argument_parser = argparse.ArgumentParser(description='自动完成疫情填报')
    argument_parser.add_argument('-r', '--reconfigure',
                                 action='store_true', help='重新配置')
    argument_parser.add_argument('-c', '--config',
                                 default=os.path.join(
                                     os.path.split(__file__)[0],
                                     'yqtb_config.ini'),
                                 help='配置文件路径')
    argument_parser.add_argument('-C', '--cookies',
                                 default=os.path.join(
                                     os.path.split(__file__)[0],
                                     'cookies.bin'),
                                 help='Cookies文件路径')
    argument_parser.add_argument('-k', '--key',
                                 default=os.path.join(
                                     os.path.split(__file__)[0],
                                     'publicKey.pem'),
                                 help='公钥文件路径')
    argument_parser.add_argument('-a', '--area',
                                 default=os.path.join(
                                     os.path.split(__file__)[0],
                                     'area_data.json'),
                                 help='地区数据文件路径')
    args = argument_parser.parse_args()
    if not os.path.exists(args.config) or args.reconfigure:
        print('开始配置')
        if os.path.exists(args.cookies):
            os.remove(args.cookies)
        return generate_config(args.config, args.key, args.area)
    error_count = 0
    yqtb = Yqtb(args.config, args.cookies)
    while True:
        try:
            yqtb.run()
        except (IndexError, FilloutFailedException):
            if error_count < 5:
                error_count += 1
                print('填报失败，正在重试：第%d次' % error_count, file=sys.stderr)
                continue
            print('已连续5次填报失败，请尝试检查配置或者手动填报。', file=sys.stderr)
            sys.exit(1)
        break


class Yqtb:
    def __init__(self, config_file: str, cookies_file: str) -> None:
        self.username = None
        self.password = None
        self.area_name = None
        self.area_code = None
        self.config_file = config_file
        self.cookies_file = cookies_file
        self.read_config()
        # 创建会话，加载Cookies
        self.session = requests.Session()
        if os.path.exists(cookies_file):
            with open(cookies_file, 'rb') as f:
                self.session.cookies.update(pickle.load(f))
        self.host = 'https://yqtb.nwpu.edu.cn'

    def run(self) -> None:
        """
        执行填报逻辑
        """
        home = self.session.get(self.host)
        # 需要登录时执行登录逻辑
        if home.url.startswith('https://uis.nwpu.edu.cn'):
            print('正在登录...')
            home = self.login(home)
        if not home.url.startswith(self.host):
            raise RuntimeError('登录失败！')
        self.submit(home.url)
        self.save_cookies()

    def save_cookies(self) -> None:
        """
        保存Cookies到文件
        """
        with open(self.cookies_file, 'wb') as f:
            pickle.dump(self.session.cookies, f)

    def login(self, res: requests.Response) -> requests.Response:
        """
        登录

        :param res: 被重定向的响应，用于获取表单参数
        :return: 登录响应
        """
        data = {'username': self.username, 'password': self.password,
                'currentMenu': 1, '_eventId': 'submit',
                'geolocation': '', 'submit': 'submit',
                'execution': re.findall(
                    r'name="execution" value="(.{4})"', res.text)[0]
                }
        return self.session.post(res.url, data=data)

    def submit(self, referer: str) -> None:
        """
        提交疫情填报表单

        :param referer: `Referer` 请求头
        :raises FilloutFailedException: 填报失败时抛出的异常
        :raises IndexError: `user_type`解析失败时抛出的异常
        """
        url = self.host + '/wx/ry/jrsb.jsp'
        headers = {'Referer': referer}
        res = self.session.get(url, headers=headers)
        params = re.findall(r'url:\'ry_util.jsp(\S*)\',', res.text)
        params = params[0] if len(params) > 0 else ''
        post_url = self.host + '/wx/ry/ry_util.jsp' + params
        name = re.findall(r',userName:\'([^\']{2,20})\'', res.text)
        name = name[0] if len(name) > 0 else ''
        # dep = re.findall(r',xymc:\'([^\']+)\'', res.text)
        # dep = dep[0] if len(dep) > 0 else ''
        # tel = re.findall(r',xssjhm:\'([0-9]{11})\'', res.text)
        # tel = tel[0] if len(tel) > 0 else ''

        # 可能会抛出IndexError异常
        user_type = re.findall(r',userType:\'(\d)\'', res.text)[0]
        data = {'actionType': 'addRbxx',
                'tbly': 'sso',
                'userLoginId': self.username,
                'szcsmc': self.area_name,
                'szcsbm': self.area_code,
                'userType': user_type,
                'userName': name,
                # 'xymc': dep,
                # 'xssjhm': tel,
                'sfyzz': '0',
                'sfqz': '0',
                'qtqksm': '',
                'ycqksm': '',
                'hsjc': '1',
                'xasymt': '1'
                }
        headers = {'Referer': url, 'Origin': self.host}
        print('正在提交...')
        post_res = self.session.post(post_url, data=data, headers=headers)
        msg = post_res.json()
        if int(msg['state']) != 1:
            print(post_res.text)
            raise FilloutFailedException('提交失败！')
        print('提交成功。')

    def read_config(self) -> None:
        """
        读取配置文件
        """
        config_parser = configparser.ConfigParser()
        config_parser.read(self.config_file)
        self.username = config_parser.get('User', 'username')
        self.password = config_parser.get('User', 'password')
        self.area_name = config_parser.get('User', 'area_name')
        self.area_code = config_parser.get('User', 'area_code')


def generate_config(config_file: str, key_file: str, area_file: str) -> None:
    """
    提示用户获取基本信息，并创建配置文件

    :param config_file: 配置文件存储路径
    :param key_file: 公钥文件路径
    :param area_file: 地区数据文件路径
    """
    username = prompt('学号：', validator=UsernameValidator())
    password = prompt('密码：', is_password=True)
    encrypted_password = encrypt_password(password, key_file)
    area_name, area_code = prompt_area(area_file)
    config_parser = configparser.ConfigParser()
    config_parser.add_section('User')
    config_parser.set('User', 'username', username)
    config_parser.set('User', 'password', encrypted_password)
    config_parser.set('User', 'area_name', area_name)
    config_parser.set('User', 'area_code', area_code)
    with open(config_file, 'w') as f:
        config_parser.write(f)
    print('成功生成配置文件')


def encrypt_password(password: str, key_file: str) -> str:
    """
    使用翱翔门户公钥加密明文密码，防止密码暴露
    公钥链接：https://uis.nwpu.edu.cn/cas/jwt/publicKey

    :param password: 明文密码
    :param key_file: 公钥文件路径
    :return: 加密后的密码
    """
    if not os.path.exists(key_file):
        raise FileNotFoundError('公钥文件不存在：%s' % key_file)
    public_key = RSA.importKey(open(key_file).read())
    cipher = PKCS1_v1_5.new(public_key)
    return '__RSA__' + base64.b64encode(
        cipher.encrypt(password.encode('utf-8'))
    ).decode('utf-8')


def prompt_area(area_file: str) -> Tuple[str, str]:
    """
    通过自动补全获取当前所在地区名称和代码

    :param area_file: 地区数据文件路径
    :return: 所在地区名称和代码
    """
    if not os.path.exists(area_file):
        raise FileNotFoundError('area_data.json文件不存在，请重新克隆本项目。')
    with open(area_file, encoding='utf-8') as f:
        area_data: Dict = json.load(f)
    area_data['在学校'] = '1'
    area_completion = copy.deepcopy(area_data)

    def set_none(dic: Dict) -> None:
        """
        将嵌套字典的非字典型节点（尾节点）设为 `None` ，以表示嵌套补全结束

        :param dic: 嵌套字典
        """
        dicts = []
        current_dict: Dict = dic
        while True:
            for k in current_dict:
                v = current_dict[k]
                if type(v) == dict:
                    dicts.append(v)
                else:
                    current_dict[k] = None
            try:
                current_dict = dicts.pop()
            except IndexError:
                return
    set_none(area_completion)
    completer = NestedCompleter.from_nested_dict(area_completion)

    def bottom_toolbar() -> str:
        """
        底部提示

        :return: 提示文本
        """
        return '按TAB键触发自动补全，通过方向键进行选择'

    area = prompt('当前所在位置：',
                  completer=completer,
                  complete_while_typing=True,
                  validator=AreaValidator(area_completion),
                  validate_while_typing=True,
                  bottom_toolbar=bottom_toolbar).strip()
    # 在学校
    if area in area_data:
        return area, area_data[area]
    province, city, district = area.split()
    return area.replace(' ', ''), area_data[province][city][district]


class AreaValidator(Validator):
    """
    地区验证器，验证用户输入的地区

    :param area_completion: 地区选项补全数据
    """

    def __init__(self, area_completion: Dict) -> None:
        self.area_completion = area_completion

    def validate(self, document: Document) -> None:
        text = document.text.strip()
        keys = text.split()
        if len(keys) == 0:
            raise ValidationError(0, '请选择省份')
        province = keys[0]
        if province not in self.area_completion:
            raise ValidationError(0, '请选择正确的省份')
        # 在学校
        if self.area_completion[province] is None:
            if len(keys) == 1:
                return
            raise ValidationError(len(province), '请删除冗余输入')
        if len(keys) == 1:
            raise ValidationError(len(text), '请选择城市')
        city = keys[1]
        if city not in self.area_completion[province]:
            raise ValidationError(len(province), '请选择正确的城市')
        if len(keys) == 2:
            raise ValidationError(text.find(city) + len(city), '请选择地区/县')
        district = keys[2]
        if district not in self.area_completion[province][city]:
            raise ValidationError(text.find(city) + len(city), '请选择正确的地区/县')
        if len(keys) > 3:
            raise ValidationError(
                text.find(district) + len(district),
                '请删除冗余输入')


class UsernameValidator(Validator):
    """
    学号验证器，验证学号的合法性
    """

    def validate(self, document: Document) -> None:
        text = document.text
        if not text:
            raise ValidationError(0, '请输入学号')
        if not text.isdigit():
            for i, c in enumerate(text):
                if not c.isdigit():
                    raise ValidationError(i, '请输入数字')
        if len(text) >= 4:
            year = int(text[:4])
            if year < 1938 or year > int(time.strftime('%Y')):
                raise ValidationError(0, '学号格式错误')
        if not len(text) == 10:
            raise ValidationError(len(text), '学号应为10位数字')


class FilloutFailedException(Exception):
    """
    填报失败异常
    """
    pass


if __name__ == '__main__':
    main()
